import os
import random
import shutil

dataset_folder = "C:\\Users\\Administrator\\Desktop\\dataset"
images_folder = os.path.join(dataset_folder, "images")
labels_folder = os.path.join(dataset_folder, "labels")


os.makedirs(os.path.join(images_folder, "train"), exist_ok=True)
os.makedirs(os.path.join(images_folder, "val"), exist_ok=True)
os.makedirs(os.path.join(labels_folder, "train"), exist_ok=True)
os.makedirs(os.path.join(labels_folder, "val"), exist_ok=True)

image_files = [os.path.splitext(file)[0] for file in os.listdir(images_folder) if file.endswith(".jpg")]

num_images = len(image_files)
num_train = int(0.6 * num_images)
num_val = num_images - num_train

random.shuffle(image_files)
train_files = image_files[:num_train]
val_files = image_files[num_train:]

for filename in train_files:
    shutil.copyfile(os.path.join(images_folder, filename + ".jpg"), os.path.join(images_folder, "train", filename + ".jpg"))
    shutil.copyfile(os.path.join(labels_folder, filename + ".txt"), os.path.join(labels_folder, "train", filename + ".txt"))

for filename in val_files:
    shutil.copyfile(os.path.join(images_folder, filename + ".jpg"), os.path.join(images_folder, "val", filename + ".jpg"))
    shutil.copyfile(os.path.join(labels_folder, filename + ".txt"), os.path.join(labels_folder, "val", filename + ".txt"))

print("Dataset split completed.")
